# -*- coding: utf-8 -*-
import os
import math
import pickle
import numpy as np
from multiprocessing import Process, Pipe
from ins_gen import extract_darknet

arch = 'yolov3'
hardware = 'VC709'
# hardware = 'VCU118'

def data_transfer(index, layer, next_layer, wgt_axi_time_list):
    act_width = 8
    wgt_width = 8
    act_util_rate = 0.53 * 1352 / 1408
    wgt_util_rate = 0.53

    # this layer
    if wgt_axi_time_list[index] > 288:
        wgt_pong_size = wgt_axi_time_list[index] - 288 
        wgt_pong_cycle = wgt_pong_size * 64 / wgt_util_rate
    else:
        wgt_pong_cycle = 0 

    # next layer
    if wgt_axi_time_list[index + 1] > 288:
        wgt_ping_size = 288 
        wgt_ping_cycle = wgt_ping_size * 64 / wgt_util_rate
    else:
        wgt_ping_cycle = wgt_axi_time_list[index + 1] * 64 / wgt_util_rate

    if layer['write_in']:
        # distance_limit_in_PEcatch = 0.5 * layer['kernel'][0] + 0.5 + layer['stride'][0]
        # input_size = (layer['inputshape'][0] - distance_limit_in_PEcatch) * layer['inputshape'][1] * math.ceil(layer['inputshape'][2]/8) * 8 * act_width
        input_size = layer['inputshape'][0] * layer['inputshape'][1] * math.ceil(layer['inputshape'][2]/8) * 8 * act_width
        input_cycle = input_size / 512 / act_util_rate

        # extra_size = distance_limit_in_PEcatch * layer['inputshape'][1] * math.ceil(layer['inputshape'][2]/8) * 8 * act_width
        # extra_cycle = extra_size / 512 / act_util_rate
    else:
        input_cycle = 0
        # extra_cycle = 0

    if layer['read_out']:
        output_size = layer['outputsize_true'][0] * layer['outputsize_true'][1] * math.ceil(layer['outputsize_true'][2]/8) * 8 * act_width
        output_cycle = output_size / 512 / act_util_rate
    else:
        output_cycle = 0

    if layer['write_out']:
        # distance_limit_out_PEcatch = 2 + int(layer['mux_sel_upsample_on'])
        # res_size = (layer['outputsize_true'][0] - distance_limit_out_PEcatch) * layer['outputsize_true'][1] * math.ceil(layer['outputsize_true'][2]/8) * 8 * act_width
        res_size = layer['outputsize_true'][0] * layer['outputsize_true'][1] * math.ceil(layer['outputsize_true'][2]/8) * 8 * act_width
        res_cycle = res_size / 512 / act_util_rate

        # extra_size = distance_limit_out_PEcatch * layer['outputsize_true'][1] * math.ceil(layer['outputsize_true'][2]/8) * 8 * act_width
        # extra_cycle += extra_size / 512 / act_util_rate
    else:
        res_cycle = 0

    cy = wgt_ping_cycle + wgt_pong_cycle + input_cycle + output_cycle + res_cycle
    # print('%s, %f, %f, %f, %f' % (key, (wgt_ping_cycle + wgt_pong_cycle)/200, (input_cycle + res_cycle)/200, output_cycle/200, cy/200))
    return cy, 0

def layer_simulator(i, config, parm, act, layer_name_list, parm_key_list, act_key_list, wgt_axi_time_list, px, py, kpf, cpf, balance_mode = 'kpf', double_mac = False):
    balance_num = 2
    pipeline_cycle = 24
    # simulate_mode = 'finegrained'
    simulate_mode = 'coarsegrained'
    layer = config[layer_name_list[i]]
    if i == len(layer_name_list) - 1:
        next_layer = None
    else:
        next_layer = config[layer_name_list[i+1]]

    layer_weight = parm[parm_key_list[i]]
    input_act = act[act_key_list[i]]

    data_transfer_cycle, extra_cycle = data_transfer(i, layer, next_layer, wgt_axi_time_list)
    total_cycle=0
    
    if simulate_mode == 'finegrained':
        for Y in range(0, layer['outputshape'][1], py):
            for X in range(0, layer['outputshape'][0], px):
                for Cout in range(0, layer['outputshape'][2], kpf):                  
                    cycle_table = np.zeros([px, py, kpf])
                    for x in range(px):
                        for y in range(py):
                            cycle = np.zeros(kpf)
                            fifo = np.zeros(kpf)
                            for KY in range(layer['kernel'][1]):
                                for KX in range(layer['kernel'][0]):
                                    index_x = KX-(layer['kernel'][0]//2)+(X+x)*layer['stride'][0]
                                    index_y = KY-(layer['kernel'][0]//2)+(Y+y)*layer['stride'][0]
                                    if index_x<0 or index_x>layer['inputshape'][0]-1 or index_y<0 or index_y>layer['inputshape'][1]-1:
                                        # for padding, input is zero
                                        cycle += layer['inputshape'][2]/(balance_num*cpf)
                                    else:
                                        w = layer_weight[Cout:Cout+kpf, :, KX, KY]
                                        a = input_act[0, :, index_x, index_y]
                                        pair_index = np.multiply(a, w[:])
                                        pair_index = np.int8(pair_index!=0)
                                        if double_mac:
                                            pair_index_or = np.copy(pair_index)
                                            for j in range(1, pair_index.shape[0], 2):
                                                pair_index_or[[j, j-1], :] = pair_index_or[[j-1, j], :]

                                            # get the index table, 1 for non-zero and 0 for zero
                                            pair_index = np.bitwise_or(pair_index, pair_index_or)
                                        #######################
                                        for Cin in range(0, layer['inputshape'][2], balance_num*cpf):
                                            # num of data in fifo can be calculated as sum
                                            fifo[:layer['outputshape'][2]-Cout] += np.sum(pair_index[:, Cin:Cin+balance_num*cpf], axis=1)
                                            # if np.max(fifo) > fifo_max_len:
                                            #     fifo_max_len = np.max(fifo)
                                            #######################
                                            tmp_b = (fifo >= cpf)
                                            # fifo_emtpy += kpf - np.sum(tmp_b)
                                            fifo = tmp_b * (fifo-cpf)
                                            # fifo = fifo-4
                                            cycle += 1
                                        
                            cycle += np.ceil(fifo/cpf)
                            cycle_table[x,y,:] = cycle      
                    
                    # kpf
                    if balance_mode == 'kpf':
                        cycle_table_or = np.copy(cycle_table)
                        for j in range(3, kpf, 4):
                            cycle_table_or[:,:,[j-1, j, j-3, j-2]] = cycle_table[:,:,[j-3, j-2, j-1, j]]
                        cycle_table_balanced = (cycle_table + cycle_table_or) / 2
                    # y
                    elif balance_mode == 'ypf':
                        cycle_table_or = np.copy(cycle_table)
                        for j in range(1, py, 2):
                            cycle_table_or[:,[j, j-1],:] = cycle_table[:,[j-1, j],:]
                        # for j in range(py//2):
                        #     cycle_table_or[:,[j+7, j],:] = cycle_table[:,[j, j+7],:]
                        cycle_table_balanced = (cycle_table + cycle_table_or) / 2
                    # x
                    elif balance_mode == 'xpf':
                        # x4
                        # cycle_table_balanced = np.zeros([4, py, kpf])
                        # cycle_table_balanced[0,:,:] = np.sum(cycle_table[0:4,:,:],axis=0)/4
                        # cycle_table_balanced[1,:,:] = np.sum(cycle_table[4:8,:,:],axis=0)/4
                        # cycle_table_balanced[2,:,:] = np.sum(cycle_table[8:12,:,:],axis=0)/4
                        # cycle_table_balanced[3,:,:] = cycle_table[12,:,:]

                        # x2
                        cycle_table_balanced = np.zeros([7, py, kpf])
                        cycle_table_balanced[0,:,:] = np.sum(cycle_table[0:2,:,:],axis=0)/2
                        cycle_table_balanced[1,:,:] = np.sum(cycle_table[2:4,:,:],axis=0)/2
                        cycle_table_balanced[2,:,:] = np.sum(cycle_table[4:6,:,:],axis=0)/2
                        cycle_table_balanced[3,:,:] = np.sum(cycle_table[6:8,:,:],axis=0)/2
                        cycle_table_balanced[4,:,:] = np.sum(cycle_table[8:10,:,:],axis=0)/2
                        cycle_table_balanced[5,:,:] = np.sum(cycle_table[10:12,:,:],axis=0)/2
                        cycle_table_balanced[6,:,:] = cycle_table[12,:,:]

                        # x13
                        # cycle_table_balanced = np.sum(cycle_table, axis=0)/13
                    elif balance_mode == 'total_balance':
                        # total_balance
                        cycle_table_balanced = np.sum(cycle_table)/(px*py*kpf)
                    elif balance_mode == 'no':
                        # no balance
                        cycle_table_balanced = cycle_table

                    total_cycle += np.ceil(np.max(cycle_table_balanced))
                    total_cycle += pipeline_cycle
    
    elif simulate_mode == 'coarsegrained':
        m=1 
        for item in layer_weight.shape:
            m *= item
        z=np.sum(layer_weight == 0)
        weight_sparsity = z/m

        m=1 
        for item in input_act.shape:
            m *= item
        z=np.sum(input_act == 0)
        act_sparsity = z/m

        if double_mac:
            sparsity = (1 - act_sparsity) * (1 - weight_sparsity * weight_sparsity)
        else:
            sparsity = (1 - act_sparsity) * (1 - weight_sparsity)

        if sparsity < 0.5:
            sparsity = 0.5

        op = layer['outputshape'][0] * layer['outputshape'][1] * layer['outputshape'][2] * \
             layer['inputshape'][2] * layer['kernel'][0] * layer['kernel'][0]

        total_cycle = op * sparsity / (px * py * cpf * kpf)
        total_cycle += pipeline_cycle * layer['outputshape'][0] * layer['outputshape'][1] * layer['outputshape'][2]/(px * py * kpf)

    else:
        raise NotImplementedError 

    return max(data_transfer_cycle, total_cycle) + extra_cycle, act_sparsity, weight_sparsity


class myProcess(Process):
    def __init__(self, processID, config, parm, act, layer_name_list, parm_key_list, act_key_list, wgt_axi_time_list, child_pipe, px, py, kpf, cpf):
        super().__init__()
        self.processID = processID
        self.cycle = 0
        self.config = config
        self.parm = parm
        self.act = act
        self.layer_name_list = layer_name_list 
        self.parm_key_list = parm_key_list
        self.act_key_list = act_key_list

        self.wgt_axi_time_list = wgt_axi_time_list
        self.child_pipe = child_pipe
        self.px = px
        self.py = py
        self.kpf = kpf
        self.cpf = cpf

    def run(self):
        total_cycle, act_sparsity, weight_sparsity = layer_simulator(self.processID, self.config, self.parm, self.act, self.layer_name_list, self.parm_key_list, self.act_key_list, self.wgt_axi_time_list, self.px, self.py, self.kpf, self.cpf)
        self.child_pipe.send([self.processID, int(total_cycle), float(act_sparsity), float(weight_sparsity)])
        # print("%d, %d" % (self.processID, self.cycle))

def resource(config, px, py, kpf, cpf, bram_data_deep, bram_wgt_deep, lut_pe=2637/32, lut_ctrl=144166, lut_mem=33935, ff_pe=3558/32, ff_ctrl=35288, ff_mem=416):
    act_width = 8
    wgt_width = 8
    bram_data_width = px * py * kpf * act_width # bit
    bram_wgt_width = 512 # bit

    # if bram_data_deep < 1024:
    # true dual port
    bram_data_deep = math.ceil(bram_data_deep/1024) * 1024
    # if bram_wgt_deep < 512:
    # simple dual port
    bram_wgt_deep = math.ceil(bram_wgt_deep/512) * 512   

    print(bram_data_width * bram_data_deep * 2 / (16 * 1024))
    print(bram_wgt_width * bram_wgt_deep / (16 * 1024))

    bram = (bram_data_width * bram_data_deep * 2 + bram_wgt_width * bram_wgt_deep) / (16 * 1024)
    dsp = px * py * cpf * kpf / 2
    lut = px * py * cpf * kpf * lut_pe + lut_ctrl + lut_mem
    ff = px * py * cpf * kpf * ff_pe + ff_ctrl + ff_mem

    '''
    for key in layer_name_list:
        layer = config[key]

        wgt_size = layer['kernel'][0] * layer['kernel'][0] * math.ceil(layer['inputshape'][2]/8) * 8 * math.ceil(layer['outputshape'][2]/8) * 8 * wgt_width
        wgt_axi_time = math.ceil(wgt_size / 4096 / 8)
        wgt_axi_time_list.append(wgt_axi_time)
    '''
    return bram, dsp, lut, ff

if __name__ == '__main__':
    if arch == 'yolov3':
        model = 'cfg/yolov3_ReLU_fuse-voc.cfg'
        
    else:
        raise NotImplementedError    

    if hardware == 'VC709':
        dsp_total = 3600
        lut_total = 433200
        ff_total = 866400
        bram_total = 1470 * 2
        cpf = 4
        kpf = 8
        px = 9
        py = 9
        bram_data_deep = 1024 * 1
        bram_wgt_deep = 36 * 1024
        frequency = 200
        inputshape = [288, 288, 3]

    elif hardware == 'VCU118':
        dsp_total = 6840
        lut_total = 1182240
        ff_total = 2364480
        bram_total = 2160 * 2
        cpf = 4
        kpf = 8
        px = 13
        py = 13
        bram_data_deep = 1024 * 1
        bram_wgt_deep = 36 * 1024
        frequency = 300
        inputshape = [416, 416, 3]
    else:
        raise NotImplementedError


    config, layer_name_list = extract_darknet(model, inputshape, px, py, kpf, bram_data_deep, bram_wgt_deep)

    wgt_width = 8

    wgt_axi_time_list = []
    for key in layer_name_list:
        layer = config[key]
        # if 'Convolution' in key:
        # bit
        wgt_size = layer['kernel'][0] * layer['kernel'][0] * math.ceil(layer['inputshape'][2]/8) * 8 * math.ceil(layer['outputshape'][2]/8) * 8 * wgt_width
        wgt_axi_time = math.ceil(wgt_size / 4096 / 8)
        wgt_axi_time_list.append(wgt_axi_time)

    wgt_axi_time_list.append(0)

    
    f=open("/home/zhangjm/backup/20210321/wgt.pkl",'rb')
    parm = pickle.load(f)
    f.close()
    
    f=open("/home/zhangjm/backup/20210321/act.pkl", 'rb')
    act = pickle.load(f)
    f.close()
   
    act_key_list=[]
    for key in act.keys():
        if 'convolutional' in key:
            act_key_list.append(key)

    parm_key_list=[]
    for key in parm.keys():
        if 'weight' in key:
            parm_key_list.append(key) 

    # for key in layer_name_list:
    #     layer = config[key]
    #     print(layer['read_out'])


    bram, dsp, lut, ff = resource(config, px, py, kpf, cpf, bram_data_deep, bram_wgt_deep)
    print('bram = %d, utilize_rate = %f' % (bram, bram/bram_total))
    print('dsp = %d, utilize_rate = %f' % (dsp, dsp/dsp_total))
    print('lut = %d, utilize_rate = %f' % (lut, lut/lut_total))
    print('ff = %d, utilize_rate = %f' % (ff, ff/ff_total))

    # assert lut <= lut_total, 'for %s, lut is not enough for %d' % (hardware, lut_total)
    # assert bram <= bram_total, 'for %s, bram is not enough for %d' % (hardware, bram_total)
    # assert dsp <= dsp_total, 'for %s, dsp is not enough for %d' % (hardware, dsp_total)
    # assert ff <= ff_total, 'for %s, ff is not enough for %d' % (hardware, ff_total)

    process_list = []
    parent_pipe_list = []
    # simulator
    for i in range(len(layer_name_list)):
        parent, child = Pipe()
        process_list.append(myProcess(i, config, parm, act, layer_name_list, parm_key_list, act_key_list, wgt_axi_time_list, child, px, py, kpf, cpf))
        parent_pipe_list.append(parent)
        process_list[i].start()
    
    total_cycle = 0
    for i in range(len(layer_name_list)):
        process_list[i].join()
        message = parent_pipe_list[i].recv()
        print('%d, %d, %f, %f' % (message[0], message[1], message[2], message[3]))
        total_cycle += message[1]
    
    print('total_cycle = ', total_cycle)
    print('total_delay = ', total_cycle / frequency * 1e-3)

